from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.types import Message


class BaseMiddleware(BaseHTTPMiddleware):

    async def __call__(self, scope, receive, send):
        try:
            await super().__call__(scope, receive, send)
        except RuntimeError as exc:
            if str(exc) == 'No response returned.':
                request = Request(scope, receive=receive)
                if await request.is_disconnected():
                    return
            raise

    async def dispatch(self, request, call_next):
        raise NotImplementedError()


async def get_body_iterator(body_iterator):
    result = ''
    for section in body_iterator:
        result += section
    return result


class api_wrap:
    def __init__(self, obj):
        self._it = iter(obj)

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            value = next(self._it)
        except StopIteration:
            raise StopAsyncIteration  # 停止迭代
        return value


async def set_body(request: Request):
    receive_ = await request._receive()

    async def receive() -> Message:
        return receive_

    request._receive = receive